import os
from yacs.config import CfgNode as CN
import yaml


_C = CN()
_C.BASE = ['']

# data settings
_C.DATA = CN()
_C.DATA.DATASET = 'BraTS20'
_C.DATA.SSL_DATASET = 'BraTS20'
_C.DATA.BATCH_SIZE = 1
_C.DATA.FOLD = 0

# BraTS 2020 path
_C.DATA.SURVIVAL_PATH = 'E:\\Brats 2020\\MICCAI_BraTS2020_TrainingData\\survival_info.csv'
_C.DATA.NAME_MAP_PATH = 'E:\\Brats 2020\\MICCAI_BraTS2020_TrainingData\\name_mapping.csv'
_C.DATA.TRAIN_ROOT_DIR = 'E:\\Brats 2020\\MICCAI_BraTS2020_TrainingData'

_C.DATA.TEST_SURVIVAL_PATH = 'E:\\Brats 2020\\MICCAI_BraTS2020_TrainingData\\survival_evaluation.csv'
_C.DATA.TEST_NAME_MAP_PATH = 'E:\\Brats 2020\\MICCAI_BraTS2020_TrainingData\\name_mapping_validation_data.csv'
_C.DATA.TEST_ROOT_DIR = 'E:\\Brats 2020\\MICCAI_BraTS2020_ValidationData'

_C.DATA.VAL_SURVIVAL_INFO_PATH = '/home/qlc/dataset/brats2020/MICCAI_BraTS2020_ValidationData/survival_evaluation.csv'
_C.DATA.VAL_NAME_MAPPING_PATH = '/home/qlc/dataset/brats2020/MICCAI_BraTS2020_ValidationData/name_mapping_validation_data.csv'
_C.DATA.VAL_TRAIN_ROOT_DIR = '/home/qlc/dataset/brats2020/MICCAI_BraTS2020_ValidationData'

# BraTS 2021 path
_C.DATA.BRATS_21_ROOT = 'E:\\Brats 2020\\MICCAI_BraTS2020_ValidationData'

_C.DATA.PATH_TO_CSV = 'log/train_data.csv' # path to val dataset
_C.DATA.PATH_TO_TEST_CSV = 'log/test_data.csv' # path to val dataset
_C.DATA.PATH_TO_SSL_CSV = 'log/ssl_data.csv' # path to val dataset

# data type
_C.DATA.PAIRED = False
_C.DATA.CANNY = False
_C.DATA.NUM_WORKERS = 4
_C.DATA.FRAC = 1.0

# augment
_C.DATA.CROP = False
_C.DATA.AUGMENT = False

# model settings
_C.MODEL = CN()
_C.MODEL.NAME = 'Unet3d'
_C.MODEL.UNET_DIMS = 24
_C.MODEL.DEPTH = 6
_C.MODEL.CHANNEL_MAPS = [64, 96, 128, 192, 256, 384, 512]

_C.MODEL.PRETRAINED = False
_C.MODEL.IN_CHANNELS = 4
_C.MODEL.NUM_CLASSES = 3

_C.MIM = CN()
_C.MIM.MASK_SIZE = (32, 32, 32)
_C.MIM.MASK_RATIO = 0.5
_C.MIM.MASK_TYPE = 'block_wise'
_C.MIM.ARC_TYPE = 'standard'
_C.MIM.LOG_NAME = 'mim'
_C.MIM.PROJECT_DEPTH = 2
_C.MIM.IN_CHANNELS = 4
_C.MIM.NUM_CLASSES = 3
_C.MIM.CHANNEL_MAPS = [64, 96, 128, 192, 256, 384, 512]

# training settings
_C.TRAIN = CN()
_C.TRAIN.NUM_EPOCHS = 300

_C.TRAIN.PRE_PTH = '/home/qlc/Model/Bra/log/pretrained.pth'
_C.TRAIN.TRAIN_LOGS_PATH = '/home/qlc/Model/Bra/log/train_log.csv'
_C.TRAIN.AE_PRETRAINED_MODEL_PATH = '/home/qlc/Model/Bra/log/best_model.pth'
_C.TRAIN.LOG_NAME = ''
_C.TRAIN.ACCUM_ITER = 1

_C.TRAIN.OPTIMIZER = CN()
_C.TRAIN.OPTIMIZER.NAME = 'AdamW'
_C.TRAIN.OPTIMIZER.EPS = 1e-8
_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999)  # for adamW
_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9
_C.TRAIN.OPTIMIZER.BASE_LR = 0.0005
_C.TRAIN.OPTIMIZER.WEIGHT_DECAY = 0.01

_C.TRAIN.LR_SCHEDULER = CN()
_C.TRAIN.LR_SCHEDULER.NAME = 'CosineAnnealingWarmRestarts'
_C.TRAIN.LR_SCHEDULER.T_WARM = 40
_C.TRAIN.LR_SCHEDULER.START_LR = 0.
_C.TRAIN.LR_SCHEDULER.MIN_LR = 0.
# _C.TRAIN.LR_SCHEDULER.T_0 = 0.

_C.TRAIN.MIM = CN()
_C.TRAIN.MIM.NUM_EPOCHS = 300
_C.TRAIN.MIM.WEIGHT_DECAY = 0.01
_C.TRAIN.MIM.BASE_LR = 0.0005
_C.TRAIN.MIM.BETAS = (0.9, 0.999)  # for adamW
_C.TRAIN.MIM.LOG_NAME = ''

_C.TRAIN.MIM.LR_SCHEDULER = CN()
_C.TRAIN.MIM.LR_SCHEDULER.NAME = 'CosineAnnealingWarmRestarts'
_C.TRAIN.MIM.LR_SCHEDULER.T_WARM = 40
_C.TRAIN.MIM.LR_SCHEDULER.START_LR = 0.
_C.TRAIN.MIM.LR_SCHEDULER.MIN_LR = 0.


# misc
_C.AMP = False # mix precision training
_C.DEVICE = 'cuda'


def _update_config_from_file(config, cfg_file):
    config.defrost()
    with open(cfg_file, 'r') as infile:
        yaml_cfg = yaml.load(infile, Loader=yaml.FullLoader)
    for cfg in yaml_cfg.setdefault('BASE', ['']):
        if cfg:
            _update_config_from_file(
                config, os.path.join(os.path.dirname(cfg_file), cfg)
            )
    print('merging config from {}'.format(cfg_file))
    config.merge_from_file(cfg_file)
    config.freeze()


def update_config(config, args):
    """Update config by ArgumentParser
    Args:
        config:
        args: ArgumentParser contains options
    Return:
        config: updated config
    """
    if args.cfg:
        _update_config_from_file(config, args.cfg)
    config.defrost()
    # if args.model:
    #     config.MODEL.NAME = args.model
    # if args.epoch:
    #     config.TRAIN.NUM_EPOCHS = args.epoch
    config.freeze()
    return config


def get_config(cfg_file=None):
    """Return a clone of config or load from yaml file"""
    config = _C.clone()
    if cfg_file:
        _update_config_from_file(config, cfg_file)
    return config


if __name__ == '__main__':
    cfg_file = 'configs/b.yaml'
    config = get_config(cfg_file)
    print(type(config))


    config = get_config(cfg_file=arguments.cfg)
    config = update_config(config=config, args=arguments)
    device = config.DEVICE







